import os
import pickle

def veri(input_path, output_path):
    output = {}
    train_names = open(os.path.join(input_path, 'name_train.txt'), 'r').readlines()
    gallery_names = open(os.path.join(input_path, 'name_test.txt'), 'r').readlines()
    query_names = open(os.path.join(input_path, 'name_query.txt'), 'r').readlines()
    
    train = []
    for train_name in train_names:
        id, cam = train_name.rstrip('\n').split('_')[0:2]
        train.append({
            "filename": train_name.rstrip('\n'),
            "image_path": os.path.join(input_path, "image_train", train_name.rstrip('\n')),
            "id": id,
            "cam": cam[1:]
        })
    output["train"] = train

    gallery = []
    for gallery_name in gallery_names:
        id, cam = gallery_name.rstrip('\n').split('_')[0:2]
        gallery.append({
            "filename": gallery_name.rstrip('\n'),
            "image_path": os.path.join(input_path, "image_test", gallery_name.rstrip('\n')),
            "id": id,
            "cam": cam[1:]
        })
    output["gallery"] = gallery
    
    query = []
    for query_name in query_names:
        id, cam = query_name.rstrip('\n').split('_')[0:2]
        query.append({
            "filename": query_name.rstrip('\n'),
            "image_path": os.path.join(input_path, "image_query", query_name.rstrip('\n')),
            "id": id,
            "cam": cam[1:]
        })
    output["query"] = query

    pickle.dump(output, open(output_path, 'wb'))

if __name__ == '__main__':
    import argparse
    args = argparse.ArgumentParser()
    args.add_argument('--input_path', type=str, default=r'D:\\VeRi')
    args.add_argument('--output_path', type=str, default='veri.pkl')
    args = args.parse_args()
    veri(args.input_path, args.output_path)